from PIL import Image
import numpy as np
import cv2
import os

"""
主要想实现的功能为：将图像中旋转的图像部分矫正还原为矩形。
导出结果不正确，有待修改
"""


def find_second_largest_contour(contours):
    # 对轮廓按照面积从大到小排序
    contours_sorted = sorted(contours, key=cv2.contourArea, reverse=True)

    # 返回第二大轮廓
    if len(contours_sorted) > 1:
        return contours_sorted[1]
    else:
        return None

def rotation_correction(image):

    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 二值化
    ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)

    # 查找轮廓
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 找到第二大轮廓
    second_largest_contour = find_second_largest_contour(contours)

    # 计算最小外接矩形
    rect = cv2.minAreaRect(second_largest_contour)
    box = cv2.boxPoints(rect)
    box = np.int0(box)

    # 创建透视变换矩阵
    width, height = image.shape[1], image.shape[0]
    src_pts = box.astype("float32")
    dst_pts = np.array([[0, height-1],
                       [0, 0],
                       [width-1, 0],
                       [width-1, height-1]], dtype="float32")
    matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)

    # 应用透视变换
    result = cv2.warpPerspective(image, matrix, (width, height))

    # 显示结果
    cv2.imshow("Original", image)
    cv2.imshow("Corrected", result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    image_base_path = 'E:\datasets\satgeoloc_dataset\query'
    image_name = '1_1016_2127_1516_2527.TIF'
    image_path = os.path.join(image_base_path, image_name)

    # 读取图像
    image = cv2.imread(image_path)
    rotation_correction(image)